package require TclOO

source /home/yi/simulation/codes/TMD/TMD.tcl

namespace import oo::*

define Charmm mixin threaded

class create ReplicaSim {
    constructor {} {
	my variable Dimensions
	my variable myname
	set Dimensions {}
	set myname [lindex [ regexp -all -inline -- {[^:]*} [self]] end]
    }

}

# core
define ReplicaSim method startNodes { {directory ""} } {
    my variable Dimensions
    my variable Values
    my variable Nodes
    my variable myname
    my variable window
    my variable Related
    my variable path
    my variable confs
    if ![info exists confs] { set confs "bin imm1pb"}
    if { [string length $directory] > 0} { cd $directory}
    
    # node list save into object_name.nodes
    set nodefile [open $myname.nodes w]
    set title "\##node "
    # prepare file title for node list
    foreach dim $Dimensions {
	append title " $dim"
	if [info exists Related($dim)] {set Rs $Related($dim)} else {set Rs {}}
	foreach relatedVar  $Rs {
	    append title " $relatedVar\($dim\)"
	}
    }
    append title " path"
    puts $nodefile $title
# prepare data for nodes.
    set Nodes [combination Values]
    dict for {id info } $Nodes {
	set data " $id "
	foreach dim $Dimensions {
	    append data " [dict get $info $dim] " 
	    # dict set Nodes $id win($dim) $window($dim,[dict get $info $dim])
	    if [info exists Related($dim)] {set Rs $Related($dim)} else {set Rs {}}
	    foreach relatedVar  $Rs {
		#puts [my varname $relatedVar]
		upvar [my varname $relatedVar] Var
		append data " $Var([dict get $info $dim])"
		dict set Nodes $id $relatedVar $Var([dict get $info $dim])
	    }
	}
	#do not put path into Node information so that can be used in replica.
	#dict set Nodes $id path [pwd]/$id
	set path($id) [pwd]/$id
	append data " [pwd]/$id "
	Charmm create $id
	$id setup [ subst {
	    $confs 
	    dir $id
	    cmd $myname.run
	    output $myname.run.out 
	    compressOutput no}]
        $id start 
        $id finished start 
        puts "started $id"
	puts $nodefile $data
    }
    close $nodefile
}

# label the begining of a simulation section. Section name will affect the output files.
define ReplicaSim method begin { section } {
     my variable Nodes
     my variable currentSection
     my variable previousSection
     my variable currentStep
     if [info exists currentSection] { set previousSection $currentSection} else {set previousSection $section}
     set currentSection $section
     set currentStep 0
     foreach {id info} $Nodes {
       $id sendCmd "label begin$section"
     }
     foreach {id info} $Nodes {
       $id finished sendCmd
     }
}
define ReplicaSim method do { cmds {until done } { condition {}} } {
    my variable Nodes
    my variable Dimensions
    set out ""
    set methods [info class methods [self class]]
    set variables [info object vars [self] ]
    dict for { node info } $Nodes {
        dict with info {
	    set nodecmds [subst $cmds]
        }
        puts "$node start: $nodecmds..."
	$node do $nodecmds $until $condition
    }
    dict for { node info } $Nodes {
	$node finished do
	puts "$node finished"
    }
}

# setup
define ReplicaSim method hasDimension { dim } {
    my variable Dimensions
    lappend Dimensions $dim
}
define ReplicaSim method setDimension { dim values } {
    my variable Dimensions
    my variable Values
    set Values($dim) $values
}
define ReplicaSim method addRelatedVarTo { dim varname {values {}}} {
    my variable Values
    my variable $varname
    my variable Related
    if ![info exists Related($dim)] {  set Related($dim) [list $varname]} else { lappend Related($dim) $varname }
    if {[llength $values] > 0} { my setRelatedVar $varname $values}
}
define ReplicaSim method setRelatedVar { varname values} {
    my variable Values
    set dim [my depend $varname]
    upvar [my varname $varname] Var
    for {set i 0} {$i <[llength $Values($dim)]} {incr i} {
	set Var([lindex $Values($dim) $i]) [lindex $values $i]
    }
}
define ReplicaSim method depend { varname } {
    my variable Related
   foreach dim [array name Related] {
    if {[lsearch $Related($dim) $varname] >= 0 } {
    return $dim 
    }
   }
   return ""
}

define ReplicaSim method copyObj objs {
    my variable Nodes
    foreach { node info}  $Nodes {
        foreach obj $objs {
	    $node copyObj $obj   
	    #puts [$node existsObj $obj]
        }
    }
}

define ReplicaSim method setup { configurations} {
my variable confs 
set confs $configurations
}
# helper

proc combination { arr } {
    upvar $arr Arr
    set name [lindex [array names Arr] 0]
    array set newArr [array get Arr]
    array unset newArr $name
    if { [llength [array names newArr]] > 0 } {
	set comb [combination newArr]
    } else {
	set id 1
	foreach value $Arr($name) {
	    dict set comb node$id $name $value
	    incr id
	}
	return $comb
    }
    set i 1
    foreach value $Arr($name) {
	foreach { id  info} $comb {
	    dict set newcomb node$i $name $value
	    foreach key [dict keys $info] {
		dict set newcomb node$i $key [dict get $info $key]
	    }
	    incr i
	}
    }
    return $newcomb
}
proc getWinWidth { values } {
    set w {}
    for {set i 0 } { $i < [llength $values]} { incr i } {
    	set center [expr double([lindex $values $i])]
    	if { $i == 0 } {
	  set left [expr double([lindex $values 1])]
	} else {
	  set left [expr double([lindex $values [expr $i-1]])]
	}
	if { $i == [expr [llength $values]-1]} { 
	  set right [expr double([lindex $values end-1])] 
	} else {
	  set right [expr double([lindex $values [expr $i+1] ] ) ] 
	}
	set width [expr abs($left-$center)/2+abs($right-$center)/2]
	lappend w $width
    }
    return $w
}

define ReplicaSim method run { times setup exchange dim when condition } {
    my variable Nodes
    my variable myname
    my variable path
    my variable Values
    my variable currentSection
    my variable currentStep
    my variable previousStep
    upvar $dim DimVar
    
    #expose variable values to condition check
    foreach {id info} $Nodes {
       eval "set $dim\($id\) [dict get $info $dim]"
    }
    # if section undefined, use object name
    if ![info exists currentSection] { my begin $myname} 
    # create a label that can be recognized by Replcia Analysis script
    my do {sendCmd "label beginReplicaSim"}
    # if this is a new section, do not read restart file. Otherwise, read restartfile from the previous step of the same section
    if {$currentStep == 0 } { set outputfiles {}} else { set outputfiles [list continueFile $currentSection[expr $currentStep-1].res]}
    #start replica steps
    for {set i $currentStep} { $i <$times +$currentStep} {incr i} {
	lappend outputfiles restartFile 
	lappend outputfiles $currentSection$i.res 
	lappend outputfiles dcdFile 
	lappend outputfiles $currentSection$i.dcd 
	
	# run ReplicaSim step
	puts "runing cycle $i"
	my do " run $setup \{ $outputfiles\} " 
	set outputfiles {}
	#get energy
	foreach {id info} $Nodes {
	    $id getEnergy 
	}
	#set energy array to be energy(node,type)
	foreach {id info} $Nodes {
	    array unset en
	    array set en [$id finished getEnergy]
	    foreach name [array names en] {
		set energy($id,$name) $en($name)
	    }
	}
	#select only neighboring nodes with selected dimension different from each other
	foreach { id info} [dict filter $Nodes value "*$dim [lindex $Values($dim) 0]*" ] {
#	puts $info
	    set info [string map [list "$dim [lindex $Values($dim) 0]" "$dim \*"] $info]
#	    puts $info
	    set nodes [dict filter $Nodes value $info]
#	    puts $nodes
	    for { set j [expr $i %2]} { $j+1 < [llength $Values($dim)]} {incr j 2} {
		set d2 [dict filter $nodes value "*$dim [lindex $Values($dim) [expr $j +1]]*"]
		set d1 [dict filter $nodes value "*$dim [lindex $Values($dim) $j]*"]
		set node1 [lindex $d1 0]
		set node2 [lindex $d2 0]
		# expose variables to evaluation
		
		#evaluate condition
#		puts $condition
#		catch {eval $condition} result 
#		puts $result
#		continue
		set doexchange [eval $condition]
		puts $doexchange
		if $doexchange {
		    puts "exchange $node1 $node2"
		    #problem? if node1,node2,node3, should node1 exchange with node2 or node3 exchange with node2?
		    #if exchang happen between node1 and node2, node2 is skipped for further exchanging.
#		    incr j
		    
		    if [file exists $path($node1)/previous] { exec rm $path($node1)/previous}
		    exec ln -s $path($node2) $path($node1)/previous
		    $node1 do [subst { $setup configure { continueFile previous/$currentSection$i.res}}]
		    if [file exists $path($node2)/previous] { exec rm $path($node2)/previous}
		    exec ln -s $path($node1) $path($node2)/previous
		    $node2 do [subst { $setup configure { continueFile previous/$currentSection$i.res}}]
               
		} else {
		   $node1 do [subst { $setup configure { continueFile $currentSection$i.res}}]
		   $node2 do [subst { $setup configure { continueFile $currentSection$i.res}}]
		}
		
	    }
	}
	
    }
    set currentStep $i
}

proc select {dict specialkey do } {
    foreach { id info} [dict filter $Nodes value "*$dim [lindex $Values($dim) 0]*" ] {
	set info [string map [list "$dim [lindex $Values($dim) 0]" "$dim \*"] $info]
       #puts $info
	set nodes [dict filter $Nodes value $info]
	for { set j 1} { $j < [llength $Values($dim)]} {incr j} {
	    set d1 [dict filter $nodes value "*$dim [lindex $Values($dim) [expr $j -1]]*"]
	    set d2 [dict filter $nodes value "*$dim [lindex $Values($dim) $j]*"]
	    set node1 [lindex $d1 0]
	    set node2 [lindex $d2 0]
	    puts "exchange $node1 $d1 $d2"
	}
    }
}
